import numpy as np
import sys
import time

# 1. 生成向量数据
def generate_vectors(num_vectors=500, dimensions=5):
    """ 
    生成向量数据
    :param num_vectors: 向量数量
    :param dimensions: 向量维度
    :return: 随机生成的向量数据库
    """ 
    np.random.seed(42)
    return np.random.rand(num_vectors, dimensions)

data_vectors = generate_vectors(num_vectors=500, dimensions=5)

# 2. LSH类定义
class LSH:
    """ 
    基于余弦相似度的LSH实现，支持性能分析
    """ 
    def __init__(self, dimensions, num_hashes, num_tables):
        self.dimensions = dimensions
        self.num_hashes = num_hashes
        self.num_tables = num_tables
        self.hash_planes = [np.random.randn(num_hashes, dimensions) for _ in range(num_tables)]
        self.tables = [{} for _ in range(num_tables)]  # 修正：使用列表推导式创建字典列表

    def hash_function(self, vector, planes):
        projections = np.dot(planes, vector)
        return ''.join(['1' if p > 0 else '0' for p in projections])

    def insert(self, vectors):
        for table_id, planes in enumerate(self.hash_planes):
            for idx, vector in enumerate(vectors):
                hash_value = self.hash_function(vector, planes)
                if hash_value not in self.tables[table_id]:
                    self.tables[table_id][hash_value] = []
                self.tables[table_id][hash_value].append(idx)

    def query(self, query_vector):
        candidates = set()
        for table_id, planes in enumerate(self.hash_planes):
            hash_value = self.hash_function(query_vector, planes)
            if hash_value in self.tables[table_id]:
                candidates.update(self.tables[table_id][hash_value])
        return list(candidates)

# 3. 初始化LSH并进行插入
num_hashes = 10
num_tables = 5
lsh = LSH(dimensions=5, num_hashes=num_hashes, num_tables=num_tables)

start_time = time.time()
lsh.insert(data_vectors)
insertion_time = time.time() - start_time

# 4. 查询向量并测量性能
query_vector = np.array([0.5, 0.5, 0.5, 0.5, 0.5])  # 修正：维度改为5，与数据向量维度一致
start_time = time.time()
result_indices = lsh.query(query_vector)
query_time = time.time() - start_time

# 5. 内存占用估算
def estimate_memory_usage(lsh):
    """ 
    估算LSH内存占用
    :param lsh: LSH 实例
    :return: 内存占用（字节）
    """ 
    memory_usage = 0
    for table in lsh.tables:
        for bucket, indices in table.items():
            memory_usage += sys.getsizeof(bucket)  # 哈希值存储大小
            memory_usage += sys.getsizeof(indices)  # 索引列表大小
            memory_usage += len(indices) * sys.getsizeof(int)  # 每个索引的大小
    return memory_usage

memory_usage = estimate_memory_usage(lsh)

# 6. 输出结果
print(f"LSH构建时间: {insertion_time:.6f} 秒")
print(f"LSH查询时间: {query_time:.6f} 秒")
print(f"内存占用估算: {memory_usage / 1024:.2f} KB")
print(f"查询向量: {query_vector}")
print(f"匹配向量索引: {result_indices}")
print(f"匹配向量数量: {len(result_indices)}")